module topo_mod

  use fiona
  use flogger
  use spark
  use params_mod
  use geog_data_mod
  use recon_mod
  use operators_mod

  implicit none

  private

  public topo_read
  public topo_regrid
  public topo_deriv
  public topo_final

  integer topo_nlon
  integer topo_nlat
  real(r_kind) topo_dlon, topo_dlat
  real(r_kind), allocatable, dimension(:    ) :: topo_lon
  real(r_kind), allocatable, dimension(:    ) :: topo_lat
  real(r_kind), allocatable, dimension(:,:  ) :: topo_zs
  real(r_kind), allocatable, dimension(:,:  ) :: topo_landfrac

contains

  subroutine topo_read()

    real(r_kind), allocatable :: tmp(:)
    integer j

    call topo_final

    if (proc%is_root()) call log_notice('Use ' // trim(topo_file) // ' as topography.')

    select case (topo_type)
    case ('etopo1')
      if (planet /= 'earth') call log_error('Topography file ' // trim(topo_file) // ' is used for the Earth!')
      call fiona_open_dataset('topo', file_path=topo_file)
      call fiona_get_dim('topo', 'x', size=topo_nlon)
      call fiona_get_dim('topo', 'y', size=topo_nlat)

      allocate(topo_lon(topo_nlon))
      allocate(topo_lat(topo_nlat))
      allocate(topo_zs (topo_nlon,topo_nlat))

      call fiona_start_input('topo')
      call fiona_input('topo', 'x', topo_lon)
      call fiona_input('topo', 'y', topo_lat)
      call fiona_input('topo', 'z', topo_zs )
      call fiona_end_input('topo')

      topo_lon = topo_lon * rad
      topo_lat = topo_lat * rad

      ! Change longitude to [0,360].
      allocate(tmp(int(topo_nlon/2)))
      tmp = topo_lon(int(topo_nlon/2)+1:)
      topo_lon(int(topo_nlon/2)+1:) = topo_lon(1:int(topo_nlon/2)) + pi2
      topo_lon(1:int(topo_nlon/2)) = tmp
      do j = 1, topo_nlat
        tmp = topo_zs(int(topo_nlon/2)+1:,j)
        topo_zs(int(topo_nlon/2)+1:,j) = topo_zs(1:int(topo_nlon/2),j)
        topo_zs(1:int(topo_nlon/2),j) = tmp
      end do
      deallocate(tmp)
    case ('gmted2010')
      if (planet /= 'earth') call log_error('Topography file ' // trim(topo_file) // ' is used for the Earth!')
      call fiona_open_dataset('topo', file_path=topo_file)
      call fiona_get_dim('topo', 'lon', size=topo_nlon)
      call fiona_get_dim('topo', 'lat', size=topo_nlat)

      allocate(topo_lon      (topo_nlon          ))
      allocate(topo_lat      (          topo_nlat))
      allocate(topo_zs       (topo_nlon,topo_nlat))
      allocate(topo_landfrac (topo_nlon,topo_nlat))

      call fiona_start_input('topo')
      call fiona_input('topo', 'lon'      , topo_lon      )
      call fiona_input('topo', 'lat'      , topo_lat      )
      call fiona_input('topo', 'htopo'    , topo_zs       )
      call fiona_input('topo', 'landfract', topo_landfrac )
      call fiona_end_input('topo')

      topo_lon = topo_lon * rad
      topo_lat = topo_lat * rad

      ! Change longitude to [0,360].
      allocate(tmp(int(topo_nlon/2)))
      tmp = topo_lon(int(topo_nlon/2)+1:)
      topo_lon(int(topo_nlon/2)+1:) = topo_lon(1:int(topo_nlon/2)) + pi2
      topo_lon(1:int(topo_nlon/2)) = tmp
      do j = 1, topo_nlat
        ! zs
        tmp = topo_zs(int(topo_nlon/2)+1:,j)
        topo_zs(int(topo_nlon/2)+1:,j) = topo_zs(1:int(topo_nlon/2),j)
        topo_zs(1:int(topo_nlon/2),j) = tmp
        ! landfrac
        tmp = topo_landfrac(int(topo_nlon/2)+1:,j)
        topo_landfrac(int(topo_nlon/2)+1:,j) = topo_landfrac(1:int(topo_nlon/2),j)
        topo_landfrac(1:int(topo_nlon/2),j) = tmp
      end do
      deallocate(tmp)
    case ('mola32')
      if (planet /= 'mars') call log_error('Topography file ' // trim(topo_file) // ' is used for the Mars!')
      call fiona_open_dataset('topo', file_path=topo_file)
      call fiona_get_dim('topo', 'longitude', size=topo_nlon)
      call fiona_get_dim('topo', 'latitude' , size=topo_nlat)

      allocate(topo_lon(topo_nlon))
      allocate(topo_lat(topo_nlat))
      allocate(topo_zs (topo_nlon,topo_nlat))

      call fiona_start_input('topo')
      call fiona_input('topo', 'longitude', topo_lon)
      call fiona_input('topo', 'latitude' , topo_lat)
      call fiona_input('topo', 'alt'      , topo_zs )
      call fiona_end_input('topo')

      ! Reverse latitude dimension from the South Pole to the North Pole.
      topo_lat = topo_lat(topo_nlat:1:-1)
      topo_zs  = topo_zs (:,topo_nlat:1:-1)
    case default
      call log_error('Unknown topo_type "' // trim(topo_type) // '"!', pid=proc%id)
    end select

  end subroutine topo_read

  subroutine topo_regrid(mesh, geog_data)

    type(cubed_sphere_mesh_type), intent(in) :: mesh
    type(geog_data_type), intent(inout) :: geog_data

    integer idom, i, j, ig, jg, n
    real(r_kind) min_lon, max_lon, min_lat, max_lat
    real(r_kind) x1(3), x2(3), x3(3), x4(3), x(3)
    integer is1, ie1, is2, ie2, js, je

    call log_notice('Calculating averaged surface height and land fraction on model grids.', pid=proc%id)
    topo_dlon = topo_lon(2) - topo_lon(1)
    topo_dlat = topo_lat(2) - topo_lat(1)
    do idom = 1, mesh%ndom
      if (.not. mesh%domain(idom)%initialized) cycle
      associate (dom      => mesh%domain(idom),               &
                 zs       => geog_data%domain(idom)%zs,       &
                 landfrac => geog_data%domain(idom)%landfrac, &
                 slmsk    => geog_data%domain(idom)%slmsk)
      !!$OMP PARALLEL DO PRIVATE(ig,min_lon,max_lon,min_lat,max_lat,x1,x2,x3,x4,x,is1,ie1,is2,ie2,js,je,n,i,j) COLLAPSE(2)
      do jg = dom%jms, dom%jme
        do ig = dom%ims, dom%ime
          min_lon = minval(dom%lon(2:5,ig,jg))
          max_lon = maxval(dom%lon(2:5,ig,jg))
          min_lat = minval(dom%lat(2:5,ig,jg))
          max_lat = maxval(dom%lat(2:5,ig,jg))

          call latlon_to_cart(dom%lon(2,ig,jg), dom%lat(2,ig,jg), x1(1), x1(2), x1(3))
          call latlon_to_cart(dom%lon(3,ig,jg), dom%lat(3,ig,jg), x2(1), x2(2), x2(3))
          call latlon_to_cart(dom%lon(4,ig,jg), dom%lat(4,ig,jg), x3(1), x3(2), x3(3))
          call latlon_to_cart(dom%lon(5,ig,jg), dom%lat(5,ig,jg), x4(1), x4(2), x4(3))

          if (max_lon - min_lon > 300 * rad) then
            is1 = int(max_lon / topo_dlon) + 1
            ie1 = topo_nlon
            is2 = 1
            ie2 = int(min_lon / topo_dlon) + 1
            if (ie2 == 1) ie2 = 0 ! ETOPO1 has 0 and 360 longitude, so drop one.
          else
            is1 = int(min_lon / topo_dlon) + 1
            ie1 = int(max_lon / topo_dlon) + 1
            is2 = 0
            ie2 = 0
          end if
          js = int((min_lat + pi0p5) / topo_dlat) + 1
          je = int((max_lat + pi0p5) / topo_dlat) + 1

          zs(ig,jg) = 0
          if(allocated(topo_landfrac))landfrac(ig,jg) = 0
          n = 0
          do j = js, je
            do i = is1, ie1
              call latlon_to_cart(topo_lon(i), topo_lat(j), x(1), x(2), x(3))
              if (orient(x1, x2, x) == 1) cycle
              if (orient(x2, x3, x) == 1) cycle
              if (orient(x3, x4, x) == 1) cycle
              if (orient(x4, x1, x) == 1) cycle
              zs(ig,jg) = zs(ig,jg) + topo_zs(i,j)
              if(allocated(topo_landfrac))then
                if(topo_landfrac(i,j)>0.5)landfrac(ig,jg) = landfrac(ig,jg) + 1
              endif
              
              n = n + 1
            end do
            do i = is2, ie2
              call latlon_to_cart(topo_lon(i), topo_lat(j), x(1), x(2), x(3))
              if (orient(x1, x2, x) == 1) cycle
              if (orient(x2, x3, x) == 1) cycle
              if (orient(x3, x4, x) == 1) cycle
              if (orient(x4, x1, x) == 1) cycle
              zs(ig,jg) = zs(ig,jg) + topo_zs(i,j)
              if(allocated(topo_landfrac))then
                if(topo_landfrac(i,j)>0.5)landfrac(ig,jg) = landfrac(ig,jg) + 1
              endif
              
              n = n + 1
            end do
          end do
          
          zs(ig,jg) = zs(ig,jg) / n
          if(allocated(topo_landfrac))then
            landfrac(ig,jg) = landfrac(ig,jg) / n
            slmsk(ig,jg) = 1
            if(landfrac(ig,jg)<0.5)then
              zs   (ig,jg) = 0 ! Remove water grid
              slmsk(ig,jg) = 0
            endif
          endif
          
        end do
      end do
      !!$OMP END PARALLEL DO
      end associate
    end do

    !call fill_halo(geog_data%array)

  end subroutine topo_regrid
  
  subroutine topo_deriv(mesh, geog_data)
    type(cubed_sphere_mesh_type), intent(in) :: mesh
    type(geog_data_type), intent(inout) :: geog_data

    integer ims, ime, jms, jme, kms, kme
    integer ids, ide, jds, jde, kds, kde
    integer irs, ire, jrs, jre
    integer dms, dme
    integer idom, i, j
    integer ds, is, js, ig
    
    real(r_kind) dzdlon, dzdlat
    real(r_kind) dzdx, dzdy
  
    call log_notice('Calculating topo derivative.', pid=proc%id)
    do idom = 1, mesh%ndom
      if (.not. mesh%domain(idom)%initialized) cycle
      call mesh%get_params(idom=idom,                          &
                           ims=ims, ime=ime, jms=jms, jme=jme, &
                           ids=ids, ide=ide, jds=jds, jde=jde, &
                           kms=kms, kme=kme, kds=kds, kde=kde, &
                           irs=irs, ire=ire, jrs=jrs, jre=jre, &
                           dms=dms, dme=dme)
      associate (dx     => mesh%domain(idom)%dx,         &
                 dy     => mesh%domain(idom)%dy,         &
                 slmsk  => geog_data%domain(idom)%slmsk, &
                 zs     => geog_data%domain(idom)%zs,    &
                 zsC    => geog_data%domain(idom)%zsC,   &
                 zsL    => geog_data%domain(idom)%zsL,   &
                 zsR    => geog_data%domain(idom)%zsR,   &
                 zsB    => geog_data%domain(idom)%zsB,   &
                 zsT    => geog_data%domain(idom)%zsT,   &
                 zsQ    => geog_data%domain(idom)%zsQ,   &
                 dzsdx  => geog_data%domain(idom)%dzsdx, &
                 dzsdxC => geog_data%domain(idom)%dzsdxC,&
                 dzsdxL => geog_data%domain(idom)%dzsdxL,&
                 dzsdxR => geog_data%domain(idom)%dzsdxR,&
                 dzsdxB => geog_data%domain(idom)%dzsdxB,&
                 dzsdxT => geog_data%domain(idom)%dzsdxT,&
                 dzsdxQ => geog_data%domain(idom)%dzsdxQ,&
                 dzsdy  => geog_data%domain(idom)%dzsdy, &
                 dzsdyC => geog_data%domain(idom)%dzsdyC,&
                 dzsdyL => geog_data%domain(idom)%dzsdyL,&
                 dzsdyR => geog_data%domain(idom)%dzsdyR,&
                 dzsdyB => geog_data%domain(idom)%dzsdyB,&
                 dzsdyT => geog_data%domain(idom)%dzsdyT,&
                 dzsdyQ => geog_data%domain(idom)%dzsdyQ )
          
      call reconstruction_2d(weno_matrics(idom),                &
                             ims,ime,jms,jme,                   &
                             irs,ire,jrs,jre,                   &
                             zs,zsL,zsR,zsB,zsT,zsQ,            &
                             dzsdxL,dzsdxR,dzsdxB,dzsdxT,dzsdxQ,&
                             dzsdyL,dzsdyR,dzsdyB,dzsdyT,dzsdyQ,&
                             dx, dy)
      
      !!$OMP PARALLEL DO PRIVATE(i) COLLAPSE(2)
      do j = jrs, jre
        do i = irs, ire
          if(geog_data%domain(idom)%slmsk(i,j)==1)then
            if(topo_deriv_order==2)then
              ! 2nd order
              dzsdx(i,j) = ( zs(i+1,j) - zs(i-1,j) )/(2.*dx)
              dzsdy(i,j) = ( zs(i,j+1) - zs(i,j-1) )/(2.*dy)
            elseif(topo_deriv_order==4)then
              ! 4th order
              dzsdx(i,j) = ( zs(i-2,j) - 8.*zs(i-1,j) + 8.*zs(i+1,j) - zs(i+2,j) )/(12.*dx)
              dzsdy(i,j) = ( zs(i,j-2) - 8.*zs(i,j-1) + 8.*zs(i,j+1) - zs(i,j+2) )/(12.*dy)
            elseif(topo_deriv_order>4)then
              ! 6th order
              dzsdx(i,j) = (-zs(i-3,j) + 9.*zs(i-2,j) - 45.*zs(i-1,j) + 45.*zs(i+1,j) - 9.*zs(i+2,j) + zs(i+3,j) )/(60.*dx)
              dzsdy(i,j) = (-zs(i,j-3) + 9.*zs(i,j-2) - 45.*zs(i,j-1) + 45.*zs(i,j+1) - 9.*zs(i,j+2) + zs(i,j+3) )/(60.*dy)
            endif
            
            dzsdxC(i,j) = dzsdx(i,j)
            dzsdyC(i,j) = dzsdy(i,j)
          endif
        end do
      end do
      !!$OMP END PARALLEL DO
      
      !!$OMP PARALLEL DO PRIVATE(i) COLLAPSE(2)
      do j = jrs, jre
        do i = irs, ire
          if(geog_data%domain(idom)%slmsk(i,j)/=1)then
            zsC(  i,j) = 0
            zsL(:,i,j) = 0
            zsR(:,i,j) = 0
            zsB(:,i,j) = 0
            zsT(:,i,j) = 0
            zsQ(:,i,j) = 0
              
            dzsdx (  i,j) = 0
            dzsdxC(  i,j) = 0
            dzsdxL(:,i,j) = 0
            dzsdxR(:,i,j) = 0
            dzsdxB(:,i,j) = 0
            dzsdxT(:,i,j) = 0
            dzsdxQ(:,i,j) = 0
            dzsdy (  i,j) = 0
            dzsdyC(  i,j) = 0
            dzsdyL(:,i,j) = 0
            dzsdyR(:,i,j) = 0
            dzsdyB(:,i,j) = 0
            dzsdyT(:,i,j) = 0
            dzsdyQ(:,i,j) = 0
          endif
        end do
      end do
      !!$OMP END PARALLEL DO
            
      do j = jrs, jre
        do i = irs, ire
          if( mesh%domain(idom)%ghost_count(i,j) > 0 )then
            do ig = 1, mesh%domain(idom)%ghost_count(i,j)
              ds = mesh%domain(idom)%ghost(i,j)%idom
              ps = mesh%domain(idom)%ghost(i,j)%ps
              is = mesh%domain(idom)%ghost(i,j)%is
              js = mesh%domain(idom)%ghost(i,j)%js
              
              iA = mesh%domain(ds)%iA
              dzdx = geog_data%domain(ds)%dzsdxQ(ps,is,js)
              dzdy = geog_data%domain(ds)%dzsdyQ(ps,is,js)
              call convert_hor_deriv_cube_to_sph(dzdlon,dzdlat,dzdx,dzdy,iA)
              
              A = mesh%domain(idom)%A
              call convert_hor_deriv_sph_to_cube(dzdx,dzdy,dzdlon,dzdlat,A)
              
            end do
          endif
        end do
      end do
      
      end associate
    end do
    
  end subroutine topo_deriv

  subroutine topo_final()

    if (allocated(topo_lon)) deallocate(topo_lon)
    if (allocated(topo_lat)) deallocate(topo_lat)
    if (allocated(topo_zs )) deallocate(topo_zs )

  end subroutine topo_final

  subroutine convert_hor_deriv_cube_to_sph(dzdlon,dzdlat,dzdx,dzdy,iA)
    real(r_kind), intent(out) :: dzdlon
    real(r_kind), intent(out) :: dzdlat
    real(r_kind), intent(in ) :: dzdx
    real(r_kind), intent(in ) :: dzdy
    real(r_kind), intent(in ) :: iA(3,3) ! inverse A matrix
    
    dzdlon = iA(1,1) * dzdx + iA(1,2) * dzdy
    dzdlat = iA(2,1) * dzdx + iA(2,2) * dzdy
    
  end subroutine convert_hor_deriv_cube_to_sph

  subroutine convert_hor_deriv_sph_to_cube(dzdx,dzdy,dzdlon,dzdlat,A)
    real(r_kind), intent(out) :: dzdx
    real(r_kind), intent(out) :: dzdy
    real(r_kind), intent(in ) :: dzdlon
    real(r_kind), intent(in ) :: dzdlat
    real(r_kind), intent(in ) :: A(3,3) ! A matrix
    
    dzdx = A(1,1) * dzdlon + A(2,1) * dzdlat
    dzdy = A(1,2) * dzdlon + A(2,2) * dzdlat
    
  end subroutine convert_hor_deriv_sph_to_cube
  
end module topo_mod
